pip install tushare pandas numpy scikit-learn matplotlib
import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# 初始化Tushare接口（替换为你的真实Token）
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


class QuantStrategy:
    def __init__(self, stock_code, start_date, end_date):
        self.stock_code = stock_code
        self.start_date = start_date
        self.end_date = end_date
        self.data = None
        self.signals = None
        self.positions = None
        self.portfolio = None

    def load_data(self):
        """仅使用免费接口加载股票数据"""
        try:
            # 获取日线行情数据（免费接口）
            daily_df = pro.daily(
                ts_code=self.stock_code,
                start_date=self.start_date.replace('-', ''),
                end_date=self.end_date.replace('-', ''),
                fields='ts_code,trade_date,open,high,low,close,vol,amount'
            )

            if daily_df.empty:
                print(f"❌ 股票 {self.stock_code} 日线数据为空")
                return False

            # 处理日期格式并排序
            daily_df['trade_date'] = pd.to_datetime(daily_df['trade_date'])
            daily_df = daily_df.sort_values('trade_date')
            daily_df = daily_df.set_index('trade_date')

            self.data = daily_df
            # 计算缺失的重要指标
            self.data['return'] = self.data['close'].pct_change()  # 日收益率

            return True
        except Exception as e:
            print(f"数据加载失败: {e}")
            return False

    def preprocess_data(self):
        """预处理数据，计算技术指标"""
        if self.data is None or self.data.empty:
            return False

        # 计算移动平均线
        self.data['MA5'] = self.data['close'].rolling(window=5).mean()
        self.data['MA10'] = self.data['close'].rolling(window=10).mean()
        self.data['MA20'] = self.data['close'].rolling(window=20).mean()
        self.data['MA60'] = self.data['close'].rolling(window=60).mean()

        # 计算相对强弱指数(RSI)
        delta = self.data['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        self.data['RSI'] = 100 - (100 / (1 + rs))

        # 计算MACD
        exp1 = self.data['close'].ewm(span=12, adjust=False).mean()
        exp2 = self.data['close'].ewm(span=26, adjust=False).mean()
        self.data['MACD'] = exp1 - exp2
        self.data['Signal'] = self.data['MACD'].ewm(span=9, adjust=False).mean()
        self.data['Histogram'] = self.data['MACD'] - self.data['Signal']

        # 计算波动率
        self.data['Volatility'] = self.data['close'].rolling(window=20).std()

        # 计算成交量变化率
        self.data['VolumeChange'] = self.data['vol'].pct_change()

        # 计算未来收益率(用于机器学习标签)
        self.data['FutureReturn'] = self.data['return'].shift(-1)
        self.data['Label'] = np.where(self.data['FutureReturn'] > 0, 1, 0)

        # 丢弃缺失值
        self.data.dropna(inplace=True)
        return True

    def price_volume_strategy(self):
        """基于量价和K线的传统量化策略"""
        signals = pd.DataFrame(index=self.data.index)
        signals['signal'] = 0

        # 策略逻辑：结合均线交叉、RSI和成交量变化
        signals.loc[
            ((self.data['MA5'] > self.data['MA20']) & (self.data['MA5'].shift(1) <= self.data['MA20'].shift(1))) |
            ((self.data['RSI'] < 30) & (self.data['RSI'] > self.data['RSI'].shift(1))) |
            ((self.data['VolumeChange'] > 0.5) & (self.data['return'] > 0.02)),
            'signal'
        ] = 1

        signals.loc[
            ((self.data['MA5'] < self.data['MA20']) & (self.data['MA5'].shift(1) >= self.data['MA20'].shift(1))) |
            ((self.data['RSI'] > 70) & (self.data['RSI'] < self.data['RSI'].shift(1))) |
            ((self.data['VolumeChange'] < -0.3) & (self.data['return'] < -0.01)),
            'signal'
        ] = -1

        self.signals = signals
        return signals

    def financial_indicator_strategy(self):
        """简化的财务指标策略（不依赖实际财务数据）"""
        print("⚠️ 没有财务数据，运行简化版财务指标策略")
        signals = pd.DataFrame(index=self.data.index)
        signals['signal'] = 0

        # 基于价格和波动率的简化策略
        signals.loc[
            (self.data['return'] > 0.02) &
            (self.data['Volatility'] < self.data['Volatility'].rolling(20).mean()),
            'signal'
        ] = 1

        signals.loc[
            (self.data['return'] < -0.02) |
            (self.data['Volatility'] > self.data['Volatility'].rolling(20).mean() * 1.5),
            'signal'
        ] = -1

        self.signals = signals
        return signals

    def machine_learning_strategy(self):
        """基于机器学习的量化策略"""
        signals = pd.DataFrame(index=self.data.index)
        signals['signal'] = 0

        # 特征选择（不包含财务指标）
        features = [
            'MA5', 'MA10', 'MA20', 'MA60', 'RSI', 'MACD', 'Signal', 'Histogram',
            'Volatility', 'VolumeChange', 'return'
        ]

        X = self.data[features]
        y = self.data['Label']

        # 划分训练集和测试集（70%训练，30%测试）
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

        # 训练随机森林分类器
        model = RandomForestClassifier(n_estimators=100, random_state=42)
        model.fit(X_train, y_train)

        # 评估模型
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"机器学习模型准确率: {accuracy:.2f}")

        # 预测并生成交易信号
        self.data['Prediction'] = model.predict(X)
        signals.loc[self.data['Prediction'] == 1, 'signal'] = 1  # 预测上涨时买入
        signals.loc[self.data['Prediction'] == 0, 'signal'] = -1  # 预测下跌时卖出

        self.signals = signals
        return signals

    def apply_stop_loss(self, position, entry_price, current_price, stop_loss_pct=0.05):
        """应用止损策略"""
        if position == 1 and (entry_price - current_price) / entry_price > stop_loss_pct:
            return -1  # 止损卖出
        elif position == -1 and (current_price - entry_price) / entry_price > stop_loss_pct:
            return 1  # 止损买入（做空场景）
        return 0  # 继续持有

    def apply_take_profit(self, position, entry_price, current_price, take_profit_pct=0.1):
        """应用止盈策略"""
        if position == 1 and (current_price - entry_price) / entry_price > take_profit_pct:
            return -1  # 止盈卖出
        elif position == -1 and (entry_price - current_price) / entry_price > take_profit_pct:
            return 1  # 止盈买入（做空场景）
        return 0  # 继续持有

    def backtest(self, initial_capital=100000.0, stop_loss_pct=0.05, take_profit_pct=0.1):
        """回测策略"""
        if self.signals is None:
            print("请先生成交易信号")
            return None

        # 初始化仓位和资产组合
        positions = pd.DataFrame(index=self.signals.index).fillna(0.0)
        positions['Position'] = self.signals['signal']

        portfolio = pd.DataFrame(index=positions.index)
        portfolio['positions'] = positions['Position']
        portfolio['holdings'] = 0.0
        portfolio['cash'] = initial_capital
        portfolio['total'] = initial_capital

        entry_price = 0.0
        current_position = 0
        trades = []  # 记录交易

        # 执行回测
        for i in range(1, len(portfolio)):
            # 获取当前价格和信号
            current_price = self.data['close'].iloc[i]
            signal = positions['Position'].iloc[i]

            # 应用止损和止盈
            if current_position != 0:
                stop_loss_signal = self.apply_stop_loss(
                    current_position, entry_price, current_price, stop_loss_pct
                )
                take_profit_signal = self.apply_take_profit(
                    current_position, entry_price, current_price, take_profit_pct
                )

                if stop_loss_signal != 0:
                    signal = stop_loss_signal
                    trades.append({
                        'date': portfolio.index[i],
                        'type': '止损',
                        'price': current_price,
                        'position': current_position,
                        'pnl': (current_price - entry_price) / entry_price if current_position == 1 else (
                                                                                                                     entry_price - current_price) / entry_price
                    })
                elif take_profit_signal != 0:
                    signal = take_profit_signal
                    trades.append({
                        'date': portfolio.index[i],
                        'type': '止盈',
                        'price': current_price,
                        'position': current_position,
                        'pnl': (current_price - entry_price) / entry_price if current_position == 1 else (
                                                                                                                     entry_price - current_price) / entry_price
                    })

            # 执行交易
            if signal != current_position:
                # 计算交易数量（考虑手续费，假设0.1%）
                if signal == 1:  # 买入
                    shares_to_buy = int(portfolio['cash'].iloc[i - 1] * (1 - 0.001) // current_price)
                    portfolio.loc[portfolio.index[i], 'holdings'] = shares_to_buy * current_price
                    portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[
                                                                    i - 1] - shares_to_buy * current_price * (1 + 0.001)
                    entry_price = current_price
                    current_position = 1

                    trades.append({
                        'date': portfolio.index[i],
                        'type': '买入',
                        'price': current_price,
                        'position': 1
                    })

                elif signal == -1 and current_position != 0:  # 卖出
                    portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[i - 1] + \
                                                                portfolio['holdings'].iloc[i - 1] * (1 - 0.001)
                    portfolio.loc[portfolio.index[i], 'holdings'] = 0
                    current_position = 0

                    trades.append({
                        'date': portfolio.index[i],
                        'type': '卖出',
                        'price': current_price,
                        'position': -1,
                        'pnl': (current_price - entry_price) / entry_price
                    })
                else:
                    # 保持上一状态
                    portfolio.loc[portfolio.index[i], 'holdings'] = portfolio['holdings'].iloc[i - 1]
                    portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[i - 1]
            else:
                # 保持上一状态
                portfolio.loc[portfolio.index[i], 'holdings'] = portfolio['holdings'].iloc[i - 1] * (
                            current_price / self.data['close'].iloc[i - 1])
                portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[i - 1]

            # 计算总资产
            portfolio.loc[portfolio.index[i], 'total'] = portfolio['holdings'].iloc[i] + portfolio['cash'].iloc[i]

        self.portfolio = portfolio
        self.trades = pd.DataFrame(trades)
        return portfolio

    def calculate_performance(self):
        """计算回测绩效"""
        if self.portfolio is None:
            print("请先执行回测")
            return None

        # 计算每日收益率
        self.portfolio['daily_return'] = self.portfolio['total'].pct_change()

        # 计算累积收益率
        total_return = (self.portfolio['total'].iloc[-1] / self.portfolio['total'].iloc[0]) - 1

        # 计算年化收益率
        days = len(self.portfolio)
        annual_return = ((1 + total_return) ** (252 / days)) - 1

        # 计算夏普比率
        risk_free_rate = 0.03  # 假设无风险利率为3%
        sharpe_ratio = (annual_return - risk_free_rate) / (self.portfolio['daily_return'].std() * (252 ** 0.5))

        # 计算最大回撤
        self.portfolio['cum_max'] = self.portfolio['total'].cummax()
        self.portfolio['drawdown'] = (self.portfolio['total'] / self.portfolio['cum_max']) - 1
        max_drawdown = self.portfolio['drawdown'].min()

        # 计算交易统计
        if hasattr(self, 'trades') and not self.trades.empty:
            trades = self.trades[self.trades['type'].isin(['卖出', '止盈', '止损'])]
            if not trades.empty:
                win_trades = trades[trades['pnl'] > 0]
                win_rate = len(win_trades) / len(trades) if len(trades) > 0 else 0
                avg_win = win_trades['pnl'].mean() if len(win_trades) > 0 else 0
                avg_loss = trades[trades['pnl'] <= 0]['pnl'].mean() if len(trades[trades['pnl'] <= 0]) > 0 else 0
                profit_factor = abs(win_trades['pnl'].sum() / trades[trades['pnl'] <= 0]['pnl'].sum()) if \
                trades[trades['pnl'] <= 0]['pnl'].sum() != 0 else float('inf')
            else:
                win_rate, avg_win, avg_loss, profit_factor = 0, 0, 0, 0
        else:
            win_rate, avg_win, avg_loss, profit_factor = 0, 0, 0, 0

        return {
            'total_return': total_return,
            'annual_return': annual_return,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'win_rate': win_rate,
            'avg_win': avg_win,
            'avg_loss': avg_loss,
            'profit_factor': profit_factor,
            'num_trades': len(self.trades) if hasattr(self, 'trades') else 0
        }

    def plot_performance(self):
        """绘制回测结果图表"""
        if self.portfolio is None:
            print("请先执行回测")
            return

        plt.figure(figsize=(14, 10))

        # 绘制资产曲线
        plt.subplot(3, 1, 1)
        plt.plot(self.portfolio['total'])
        plt.title('Portfolio Value Over Time')
        plt.xlabel('Date')
        plt.ylabel('Portfolio Value')
        plt.grid(True)

        # 绘制回撤曲线
        plt.subplot(3, 1, 2)
        plt.fill_between(self.portfolio.index, self.portfolio['drawdown'], 0, color='red', alpha=0.3)
        plt.title('Portfolio Drawdown Over Time')
        plt.xlabel('Date')
        plt.ylabel('Drawdown')
        plt.grid(True)

        # 绘制交易点
        if hasattr(self, 'trades') and not self.trades.empty:
            plt.subplot(3, 1, 3)
            plt.plot(self.data.index, self.data['close'], label='Price')

            # 标记买入点
            buy_points = self.trades[self.trades['type'] == '买入']
            if not buy_points.empty:
                plt.scatter(
                    buy_points['date'],
                    [self.data.loc[date, 'close'] for date in buy_points['date']],
                    marker='^', color='g', s=100, label='Buy'
                )

            # 标记卖出点
            sell_points = self.trades[self.trades['type'].isin(['卖出', '止盈', '止损'])]
            if not sell_points.empty:
                plt.scatter(
                    sell_points['date'],
                    [self.data.loc[date, 'close'] for date in sell_points['date']],
                    marker='v', color='r', s=100, label='Sell'
                )

            plt.title('Trading Signals')
            plt.xlabel('Date')
            plt.ylabel('Price')
            plt.legend()
            plt.grid(True)

        plt.tight_layout()
        plt.show()


def compare_strategies(stock_list, start_date, end_date):
    """比较不同股票和策略的表现"""
    results = []

    for stock in stock_list:
        print(f"\n回测股票: {stock}")
        strategy = QuantStrategy(stock, start_date, end_date)

        if not strategy.load_data():
            continue

        if not strategy.preprocess_data():
            continue

        # 策略1: 量价策略
        print("\n测试量价策略...")
        strategy.price_volume_strategy()
        strategy.backtest()
        performance1 = strategy.calculate_performance()
        if performance1:
            results.append({
                'stock': stock,
                'strategy': 'Price-Volume',
                'total_return': performance1['total_return'],
                'annual_return': performance1['annual_return'],
                'sharpe_ratio': performance1['sharpe_ratio'],
                'max_drawdown': performance1['max_drawdown'],
                'win_rate': performance1['win_rate'],
                'num_trades': performance1['num_trades']
            })
            print(f"量价策略总收益率: {performance1['total_return']:.2%}")
            strategy.plot_performance()

        # 策略2: 简化的财务指标策略
        print("\n测试简化版财务指标策略...")
        strategy.financial_indicator_strategy()
        strategy.backtest()
        performance2 = strategy.calculate_performance()
        if performance2:
            results.append({
                'stock': stock,
                'strategy': 'Financial Indicator',
                'total_return': performance2['total_return'],
                'annual_return': performance2['annual_return'],
                'sharpe_ratio': performance2['sharpe_ratio'],
                'max_drawdown': performance2['max_drawdown'],
                'win_rate': performance2['win_rate'],
                'num_trades': performance2['num_trades']
            })
            print(f"财务指标策略总收益率: {performance2['total_return']:.2%}")
            strategy.plot_performance()

        # 策略3: 机器学习策略
        print("\n测试机器学习策略...")
        strategy.machine_learning_strategy()
        strategy.backtest()
        performance3 = strategy.calculate_performance()
        if performance3:
            results.append({
                'stock': stock,
                'strategy': 'Machine Learning',
                'total_return': performance3['total_return'],
                'annual_return': performance3['annual_return'],
                'sharpe_ratio': performance3['sharpe_ratio'],
                'max_drawdown': performance3['max_drawdown'],
                'win_rate': performance3['win_rate'],
                'num_trades': performance3['num_trades']
            })
            print(f"机器学习策略总收益率: {performance3['total_return']:.2%}")
            strategy.plot_performance()

    # 转换结果为DataFrame并展示
    if results:
        results_df = pd.DataFrame(results)
        return results_df
    else:
        print("没有成功回测任何策略")
        return None


if __name__ == "__main__":
    # 设置回测参数
    start_date = "2020-01-01"
    end_date = "2023-12-31"

    # 选择要回测的股票列表（A股示例）
    stock_list = ["000001.SZ", "600000.SH", "601318.SH", "601857.SH", "002594.SZ"]

    # 比较不同策略
    results = compare_strategies(stock_list, start_date, end_date)

    # 打印结果
    if results is not None:
        print("\n===== 策略比较结果 =====")
        print(results.sort_values(by=['stock', 'total_return'], ascending=[True, False]))

        # 找出总收益率为正的策略
        positive_results = results[results['total_return'] > 0]
        if not positive_results.empty:
            print("\n===== 总收益率为正的策略 =====")
            print(positive_results)
        else:
            print("\n⚠️ 没有策略在所有股票上获得正收益")